from cnocr import CnOcr
from pathlib import Path
import os

def extract_text_from_documents(folder_path, filtername = ''):
    ocr = CnOcr( det_model_backend = 'pytorch',  rec_model_backend = 'pytorch')
    all_text = []
    
    for filename in os.listdir(folder_path):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            isNeedUse = True
            if filtername != '':
                isNeedUse = filtername == filename

            # print(f'{filtername} isNeedUse: {filename}  {isNeedUse}')    
            
            img_path = os.path.join(folder_path, filename)
            if isNeedUse:
                print(img_path)
                try:
                    result = ocr.ocr(img_path)
                    text = ' '.join([''.join(line['text']) for line in result])
                    all_text.append({
                        'filename': filename,
                        'text': text
                    })
                    print(f"处理完成: {filename}")
                except Exception as e:
                    print(f"处理失败 {filename}: {e}")
    
    return all_text

# 使用示例
documents = extract_text_from_documents(Path('D:/projects/data/ocr'), "test2.png")
for doc in documents:
    print(f"文件: {doc['filename']}")
    print(f"内容: {doc['text']}")  # 只显示前100个字符